package de.distMLP.messages;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;

import de.mlp_distributed.mlp.math.mahout.Matrix;
import de.mlp_distributed.mlp.math.mahout.MatrixWritable;

public class WeightMatrixMessage extends HamaMessage {

	private Matrix[] matrix;
	private boolean stopTraining = false;
	private final Kryo kryo = new Kryo();

	byte[] serializedArray = null;

	private double cost = 0;

	public WeightMatrixMessage(final Matrix[] matrix, final boolean stopTraining, final double cost) throws IOException {
		this();
		this.matrix = matrix;
		this.stopTraining = stopTraining;
		this.cost = cost;
		this.serializedArray = this.serialize();
	}

	public WeightMatrixMessage() {
		super();
	}

	public byte[] serialize() throws IOException {
		final ByteArrayOutputStream out = new ByteArrayOutputStream();

		final Output fileOutput = new Output(out);
		this.kryo.writeObjectOrNull(fileOutput, this.matrix, Matrix[].class);
		this.kryo.writeObjectOrNull(fileOutput, this.stopTraining, Boolean.class);
		this.kryo.writeObjectOrNull(fileOutput, this.cost, Double.class);
		fileOutput.flush();
		fileOutput.close();
		out.flush();
		out.close();
		return out.toByteArray();
	}

	public void deserialize(final byte[] input) throws IOException {
		final ByteArrayInputStream in = new ByteArrayInputStream(input);
		final Input fileOutput = new Input(in);
		this.matrix = this.kryo.readObjectOrNull(fileOutput, Matrix[].class);
		this.stopTraining = this.kryo.readObjectOrNull(fileOutput, Boolean.class);
		this.cost = this.kryo.readObjectOrNull(fileOutput, Double.class);
	}

	@Override
	public void readFields(final DataInput in) throws IOException {
		final int size = in.readInt();
		final byte[] data = new byte[size];
		in.readFully(data);
		this.deserialize(data);
	}

	@Override
	public void write(final DataOutput out) throws IOException {
		if (this.serializedArray == null) {
			this.serializedArray = this.serialize();
		}
		out.writeInt(this.serializedArray.length);
		out.write(this.serializedArray);
	}

	protected MatrixArrayWritable createMatrixArrayWritable(final Matrix[] matrix) {
		final MatrixWritable[] list = new MatrixWritable[matrix.length];
		for (int i = 0; i < matrix.length; i++) {
			list[i] = new MatrixWritable(matrix[i]);
		}
		return new MatrixArrayWritable(list);
	}

	protected Matrix[] getMatricesFromMatrixArrayWritable(final MatrixArrayWritable mar) {
		final Matrix[] list = new Matrix[mar.get().length];
		for (int i = 0; i < mar.get().length; i++) {
			final MatrixWritable tmp = (MatrixWritable) mar.get()[i];
			list[i] = tmp.get();
		}
		return list;
	}

	public boolean isStopTraining() {
		return this.stopTraining;
	}

	public Matrix[] getMatrix() {
		return this.matrix;
	}

	public double getCost() {
		return this.cost;
	}
}
